
meters_to_inches = 39.37007874

import bpy
from math import degrees

class Format_BVH(object):
    def __init__(self, anim):
        self.anim = anim

    def export(self, filepath):
        print(self.anim, "exporting to", filepath)
        file = open(filepath, "w", encoding="utf8", newline="\n")
        try:
            file.write("HIERARCHY\n")

            if len(self.anim.armature.root_bones) == 1:
                indent = 0
                self.write_bone_recursive_nodes_bvh(self.anim.armature.root_bones[0], file, indent)
            else:
                file.write("ROOT ROOT\n")
                file.write("{\n")
                file.write("\tOFFSET 0.0 0.0 0.0\n")
                file.write("\tCHANNELS 0\n")
                key = None
                indent = 1
                for bone in self.anim.armature.root_bones:
                    self.write_bone_recursive_nodes(bone, file, indent)
                file.write("}\n")
                
            frame_count = self.anim.frames.stop - self.anim.frames.start

            file.write("MOTION\n")
            file.write("Frames: %d\n" % frame_count)
            file.write("Frame Time: %.6f\n" % self.anim.frame_duration)

            for frame in self.anim.frames:
                #self.anim.armature.update_posedata(self)

                for bone in self.anim.armature.root_bones:
                    self.recursive_write_bone_frame_bvh(bone, file, frame)

                file.write("\n")
        finally:
            print("closing file")
            file.close()
            print("closed file")

        print("BVH Exported: %s frames:%d\n" % (filepath, frame_count))
        return {'FINISHED'}

    def write_bone_recursive_nodes_bvh(self, bone, file, indent):
        if not bone.mandatory: return
        print(bone, bone.pose_bone)

        indent_str = "\t" * indent

        def loc_scaled(indent_str, loc):
            scale = meters_to_inches
            return (indent_str, loc.x * scale, loc.y * scale, loc.z * scale)

        bone.use_rot = True # because of bad Firestorm-BVH-IMPORT (rotations mandatory for given bones)

        pose_bone = bone.pose_bone
        rest_bone = bone.rest_bone
        loc = rest_bone.head

        file.write("%s%s %s\n" % (indent_str, "JOINT" if indent else "ROOT", bone.name))

        file.write("%s{\n" % indent_str)
        file.write("%s\tOFFSET %.6f %.6f %.6f\n" % loc_scaled(indent_str, loc))

        t_pos = "" if not bone.use_loc else " Xposition Yposition Zposition"
        t_rot = "" if not bone.use_rot else " Xrotation Yrotation Zrotation"
        t_cnt = 0
        if bone.use_loc: t_cnt += 3
        if bone.use_rot: t_cnt += 3

        file.write("%s\tCHANNELS %i%s%s\n" % (indent_str, t_cnt, t_pos, t_rot))

        if len(bone.children) > 0:
            for child_bone in bone.children:
                self.write_bone_recursive_nodes_bvh(child_bone, file, indent + 1)
        else:
            file.write("%s\tEnd Site\n" % indent_str)
            file.write("%s\t{\n" % indent_str)
            
            loc = bone.tail_location - bone.head_location

            file.write("%s\t\tOFFSET %.6f %.6f %.6f\n" % loc_scaled(indent_str, loc))
            file.write("%s\t}\n" % indent_str)

        file.write("%s}\n" % indent_str)

    def recursive_write_bone_frame_bvh(self, bone, file, frame):
        loc, rot = bone.frame_values[frame]
        
        if bone.use_loc:
            file.write("%.6f %.6f %.6f " % (loc * meters_to_inches)[:])

        if bone.use_rot:
            file.write("%.6f %.6f %.6f " % tuple(-degrees(val) for val in rot.to_euler()))
        
        for child_bone in bone.children:
            self.recursive_write_bone_frame_bvh(child_bone, file, frame)

def register():
    print("registering: OSSL Format BVH")
    print("registered: OSSL Format BVH")

def unregister():
    print("unregistering: OSSL Format BVH")
    print("unregistered: OSSL Format BVH")

import ossl.util as util

if __name__ == "__main__":
    util.reset_addon()
